# PiecePicker.py
#
# Copyright (c) 2006 Michael Hobbs
#
# This file is part of iTorrent.
#
# iTorrent is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# iTorrent is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iTorrent; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA


from random import shuffle

from BitTorrent.PiecePicker import PiecePicker


class patched_PiecePicker(PiecePicker):
	def __init__(self, *args, **kwargs):
		super(patched_PiecePicker, self).__init__(*args, **kwargs)
		if '#only-pieces' not in self.config:
			return
		pieces = self.config['#only-pieces']
		self.numpieces = len(pieces)
		self.scrambled = pieces[:]
		shuffle(self.scrambled)

	def next(self, havefunc, seed = False):
		# See what the parent class would return. If the returned index is in the
		# started list, return it.
		r = super(patched_PiecePicker, self).next(havefunc, seed)
		if '#piece-picker-next' not in self.config:
			return r
		if r is None or r in self.started or r in self.seedstarted:
			return r
		picked = self.config['#piece-picker-next'](self.have)
		if picked is not None and havefunc(picked) and picked < self.numpieces:
			return picked
		return r
